#!/bin/bash
#SBATCH --nodes=14
#SBATCH --partition=wzhdnormal
#SBATCH --ntasks-per-node=32
#SBATCH --time=48:00:00
#SBATCH --job-name=equ_run_cpu
#SBATCH -e log/equ_run-%j.out
#SBATCH -o log/equ_run-%j.out

source ~/.bashrc
source /work/home/aixplorerbio_wz/miniconda3/bin/activate amber
source /work/home/aixplorerbio_wz/software/gromacs-2020.3/bin/GMXRC

if [ $# != 1 ]; then
    echo "USAGE: $0 $PROJECT_HOME"
    echo " e.g.: $0 ../"
    exit 1
fi

pwd=$(pwd)
PROJECT_HOME=$(cd $1 && pwd)
WORK_PATH=$PROJECT_HOME/work_dir/job

if [ ! -d "$WORK_PATH" ]; then
    echo "directory \"$WORK_PATH\" not exists"
    exit 1
fi

cd $WORK_PATH

echo " Start equilibration"
# function min_equ {
# [ ! -f $1/min01.tpr ] && { gmx grompp -c $1/system.pdb -r $1/system.pdb -f $1/min01.mdp -p $1/system.top -o $1/min01.tpr -n $1/index.ndx -po $1/min01_mdout.mdp -maxwarn 7 > $1/min01_grompp.log 2>&1 || { echo "Failed to run command  at line ${LINENO}" && exit; } }
# [ ! -f $1/min01.gro ] && { ( cd $1/ && gmx mdrun --deffnm min01 -nt 1 > min01_mdrun_log.log 2>&1 ) || { echo "Failed to run command  at line ${LINENO}" && exit; } }
# [ ! -f $1/min02.tpr ] && { gmx grompp -c $1/min01.gro -r $1/min01.gro -f $1/min02.mdp -p $1/system.top -o $1/min02.tpr -n $1/index.ndx -po $1/min02_mdout.mdp -maxwarn 7 > $1/min02_grompp.log 2>&1 || { echo "Failed to run command  at line ${LINENO}" && exit; } }
# [ ! -f $1/min02.gro ] && { ( cd $1/ && gmx mdrun --deffnm min02 -nt 1 > min02_mdrun_log.log 2>&1 ) || { echo "Failed to run command  at line ${LINENO}" && exit; } }
# [ ! -f $1/min03.tpr ] && { gmx grompp -c $1/min02.gro -r $1/min02.gro -f $1/min03.mdp -p $1/system.top -o $1/min03.tpr -n $1/index.ndx -po $1/min03_mdout.mdp -maxwarn 7 > $1/min03_grompp.log 2>&1 || { echo "Failed to run command  at line ${LINENO}" && exit; } }
# [ ! -f $1/min03.gro ] && { ( cd $1/ && gmx mdrun --deffnm min03 -nt 1 > min03_mdrun_log.log 2>&1 ) || { echo "Failed to run command  at line ${LINENO}" && exit; } }
# [ ! -f $1/nve.tpr ] && { gmx grompp -c $1/min03.gro -r $1/min03.gro -f $1/nve.mdp -p $1/system.top -o $1/nve.tpr -n $1/index.ndx -po $1/nve_mdout.mdp -maxwarn 7 > $1/nve_grompp.log 2>&1 || { echo "Failed to run command  at line ${LINENO}" && exit; } }
# [ ! -f $1/nve.gro ] && { ( cd $1/ && gmx mdrun --deffnm nve -nt 1 > nve_mdrun_log.log 2>&1 ) || { echo "Failed to run command  at line ${LINENO}" && exit; } }
# [ ! -f $1/nvt.tpr ] && { gmx grompp -c $1/nve.gro -r $1/nve.gro -f $1/nvt.mdp -p $1/system.top -o $1/nvt.tpr -n $1/index.ndx -po $1/nvt_mdout.mdp -maxwarn 7 > $1/nvt_grompp.log 2>&1 || { echo "Failed to run command  at line ${LINENO}" && exit; } }
# [ ! -f $1/nvt.gro ] && { ( cd $1/ && gmx mdrun --deffnm nvt -nt 1 > nvt_mdrun_log.log 2>&1 ) || { echo "Failed to run command  at line ${LINENO}" && exit; } }
# [ ! -f $1/npt.tpr ] && { gmx grompp -c $1/nvt.gro -r $1/nvt.gro -f $1/npt.mdp -p $1/system.top -o $1/npt.tpr -n $1/index.ndx -po $1/npt_mdout.mdp -maxwarn 7 > $1/npt_grompp.log 2>&1 || { echo "Failed to run command  at line ${LINENO}" && exit; } }
# [ ! -f $1/npt.gro ] && { ( cd $1/ && gmx mdrun --deffnm npt -nt 1 > npt_mdrun_log.log 2>&1 ) || { echo "Failed to run command  at line ${LINENO}" && exit; } }
# # [ ! -f $1/md_nocont.tpr ] && { gmx grompp -c $1/npt.gro -r $1/npt.gro -t $1/npt.cpt -f $1/md_nocont.mdp -po $1/mdout_nocont.mdp -o $1/md_nocont.tpr -n $1/index.ndx -p $1/system.top -maxwarn 7 > $1/mdout_nocont.log 2>&1 || { echo "Failed to run command  at line ${LINENO}" && exit; } }
# [ ! -f $1/md.tpr ] && { gmx grompp -c $1/npt.gro -r $1/npt.gro -t $1/npt.cpt -f $1/md.mdp -p $1/system.top -o $1/md.tpr -n $1/index.ndx -po $1/mdout.mdp -maxwarn 7 > $1/mdout.log 2>&1 || { echo "Failed to run command  at line ${LINENO}" && exit; } }
# }

n_jobs=0
max_jobs=90
pids=""

function waiting {
    for pid in $1; do
        wait $pid
    done
}

srun hostname | sort | uniq | awk '{printf "%s\n",$1}' >hostfile
hostlist=$(cat hostfile)
# ligands=`ls ligands/*.sdf`

if [ $(grep -c "" hostfile) -ne $(ls | egrep "_md+$" | wc -l) ]; then
    echo "Allocked node number not equal to md folders."
    exit
fi

nodes=""
nodeIdx=0
for host in $hostlist; do
    nodes[$nodeIdx]=$host
    let nodeIdx++
done

nodeIdx=0
for lig_name in $(ls -d *_md); do
    echo ${lig_name}
    echo "ssh ${nodes[$nodeIdx]} "$PROJECT_HOME/work_dir/min_equ.sh $WORK_PATH/${lig_name}" &"
    ssh ${nodes[$nodeIdx]} "$PROJECT_HOME/work_dir/min_equ.sh $WORK_PATH/${lig_name}" &
    pids="$pids $!"
    let nodeIdx++
    n_jobs=$((n_jobs + 1))
    if (($n_jobs > max_jobs)); then
        waiting $pids
        n_jobs=0
        pids=""
    fi
done
waiting $pids

echo " Finished equilibration"

echo "Start run"
if [ $(ls *_md/md.gro 2>/dev/null | wc -l) -eq $(ls | egrep "_md+$" | wc -l) ]; then
    echo "[INFO] Run step already done. Going on."
else
    #[ $(ls *_md/lambda.gro 2>/dev/null | wc -l) -ne 0 ] && { echo "[ERROR] Some of the lambda windows are complete, but not all of them. Because this run uses HREX, I cannot continue"; exit 1; }
    # mpirun -np $(ls | egrep "_md+$" | wc -l) mdrun_mpi -v --deffnm md -cpi md.cpt -multidir $(ls -v | egrep "_md+$") >&feprun_$(date "+%H%M%S_%d%m%Y").log || {
    #     echo "mdrun failed at line ${LINENO} "
    #     exit -1
    # }
    echo "start sbatch equ_run_mpi.sh..."
fi
